#include <torch/extension.h>

#include <cmath>
#include <iostream>
#include <vector>

torch::Tensor check_win_cuda(torch::Tensor board, torch::Tensor action);
torch::Tensor get_status(torch::Tensor board);

torch::Tensor check_win(torch::Tensor board,
                                     torch::Tensor action) {
  const static int dirs[4][2] = {{1, 0}, {0, 1}, {1, 1}, {1, -1}};
  auto batch_size = board.size(0);

  auto win =
      torch::zeros({batch_size}, torch::TensorOptions().dtype(torch::kBool));

  auto w = board.size(3), h = board.size(2);

  auto board_it = board.accessor<bool, 4>();
  auto win_it = win.data_ptr<bool>();
  auto action_it = action.data_ptr<long>();
  for (int i = 0; i < batch_size; i++) {
    auto j = action_it[i] / w;
    auto k = action_it[i] % w;
    board_it[i][0][j][k] = 1;
    for (auto& dir : dirs) {
      int count = 1;
      int x = j + dir[0], y = k + dir[1];
      while (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
             board_it[i][0][x][y])
        x += dir[0], y += dir[1], count++;
      x = j - dir[0], y = k - dir[1];
      while (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
             board_it[i][0][x][y])
        x -= dir[0], y -= dir[1], count++;
      if (count >= 5) {
        win_it[i] = 1;
        break;
      }
    }
    // augmentation
    // if (valid_cnt == 1 || win_it[i]) {
    //   // for (j = 0; j < board.size(2); j++) {
    //   //   for (k = 0; k < board.size(3); k++) {
    //   //     if (board_it[i][0][j][k])
    //   //       std::cerr << 'x';
    //   //     else if (board_it[i][1][j][k])
    //   //       std::cerr << 'o';
    //   //     else
    //   //       std::cerr << ' ';
    //   //   }
    //   //   std::cerr << std::endl;
    //   // }
    //   board[i] = 0;
    // } else {
    //   board[i] = torch::rot90(board[i], rand() % 4, {1, 2});
    //   if (rand() % 2) board[i] = torch::flip(board[i], {1});
    // }
  }
  return win;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("check_win", &check_win, "game forward");
  m.def("check_win_cuda", &check_win_cuda, "game forward");
  m.def("get_status", &get_status, "get_status");
}
